/**
 * Copyright (c) 2023-2024 Huawei Technologies Co., Ltd.
 *
 * ascendc_ops is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *
 * http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

/*!
 * \file flash_attention_score_grad_tiling_common.cpp
 * \brief
 */

#include "flash_attention_score_grad_tiling_common.h"
#include "log/ops_log.h"

namespace optiling {

ge::graphStatus CheckSoftmaxMaxShape(gert::TilingContext *context, uint32_t b, uint32_t n1, uint32_t s1)
{
    auto softmaxMaxShape = context->GetOptionalInputShape(SOFTMAX_MAX);
    if (softmaxMaxShape == nullptr) {
        return ge::GRAPH_SUCCESS;
    }
    auto softmaxMaxShapeDim = softmaxMaxShape->GetStorageShape().GetDimNum();
    if (softmaxMaxShapeDim != 4) { // softmaxMax only support 4 dimensions
        OPS_LOG_E(context, "The shape of softmaxMax is invalid, got %ld dimensions", softmaxMaxShapeDim);
        return ge::GRAPH_FAILED;
    }
    auto dim0 = softmaxMaxShape->GetStorageShape().GetDim(0); // 0:b
    auto dim1 = softmaxMaxShape->GetStorageShape().GetDim(1); // 1:n1
    auto dim2 = softmaxMaxShape->GetStorageShape().GetDim(2); // 2:s1
    auto dim3 = softmaxMaxShape->GetStorageShape().GetDim(3); // 3:8
    if (dim0 != b || dim1 != n1 || dim2 != s1 || dim3 != 8) { // softmaxMax pad to 8
        OPS_LOG_E(context, "The shape of softmaxMax is invalid, got (%ld,%ld,%ld,%ld), should be (%u,%u,%u,8)", dim0,
                  dim1, dim2, dim3, b, n1, s1);
        return ge::GRAPH_FAILED;
    }
    return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckTndSoftmaxMaxShape(gert::TilingContext *context, uint32_t t1, uint32_t n1)
{
    auto softmaxMaxShape = context->GetOptionalInputShape(SOFTMAX_MAX);
    if (softmaxMaxShape == nullptr) {
        return ge::GRAPH_SUCCESS;
    }
    auto softmaxMaxShapeDim = softmaxMaxShape->GetStorageShape().GetDimNum();
    if (softmaxMaxShapeDim != 3) { // TND softmaxMax only support 3 dimensions
        OPS_LOG_E(context, "The shape of softmaxMax is invalid, got %ld dimensions", softmaxMaxShapeDim);
        return ge::GRAPH_FAILED;
    }
    auto dim0 = softmaxMaxShape->GetStorageShape().GetDim(0); // 0:t1
    auto dim1 = softmaxMaxShape->GetStorageShape().GetDim(1); // 1:n1
    auto dim2 = softmaxMaxShape->GetStorageShape().GetDim(2); // 2:8
    if (dim0 != t1 || dim1 != n1 || dim2 != 8) {              // softmaxMax pad to 8
        OPS_LOG_E(context, "The shape of softmaxMax is invalid, got (%ld,%ld,%ld), should be (%u,%u,8)", dim0, dim1,
                  dim2, t1, n1);
        return ge::GRAPH_FAILED;
    }
    return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckSoftmaxSumShape(gert::TilingContext *context, uint32_t b, uint32_t n1, uint32_t s1)
{
    auto softmaxSumShape = context->GetOptionalInputShape(SOFTMAX_SUM);
    if (softmaxSumShape == nullptr) {
        return ge::GRAPH_SUCCESS;
    }
    auto softmaxSumShapeDim = softmaxSumShape->GetStorageShape().GetDimNum();
    if (softmaxSumShapeDim != 4) { // softmaxSum only support 4 dimensions
        OPS_LOG_E(context, "The shape of softmaxSum is invalid, got %ld dimensions", softmaxSumShapeDim);
        return ge::GRAPH_FAILED;
    }
    auto dim0 = softmaxSumShape->GetStorageShape().GetDim(0); // 0:b
    auto dim1 = softmaxSumShape->GetStorageShape().GetDim(1); // 1:n1
    auto dim2 = softmaxSumShape->GetStorageShape().GetDim(2); // 2:s1
    auto dim3 = softmaxSumShape->GetStorageShape().GetDim(3); // 3:8
    if (dim0 != b || dim1 != n1 || dim2 != s1 || dim3 != 8) { // softmaxSum pad to 8
        OPS_LOG_E(context, "The shape of softmaxSum is invalid, got (%ld,%ld,%ld,%ld), should be (%u,%u,%u,8)", dim0,
                  dim1, dim2, dim3, b, n1, s1);
        return ge::GRAPH_FAILED;
    }
    return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckTndSoftmaxSumShape(gert::TilingContext *context, uint32_t t1, uint32_t n1)
{
    auto softmaxSumShape = context->GetOptionalInputShape(SOFTMAX_SUM);
    if (softmaxSumShape == nullptr) {
        return ge::GRAPH_SUCCESS;
    }
    auto softmaxSumShapeDim = softmaxSumShape->GetStorageShape().GetDimNum();
    if (softmaxSumShapeDim != 3) { // TND softmaxSum only support 3 dimensions
        OPS_LOG_E(context, "The shape of softmaxSum is invalid, got %ld dimensions", softmaxSumShapeDim);
        return ge::GRAPH_FAILED;
    }
    auto dim0 = softmaxSumShape->GetStorageShape().GetDim(0); // 0:t1
    auto dim1 = softmaxSumShape->GetStorageShape().GetDim(1); // 1:n1
    auto dim2 = softmaxSumShape->GetStorageShape().GetDim(2); // 2:8
    if (dim0 != t1 || dim1 != n1 || dim2 != 8) {              // softmaxSum pad to 8
        OPS_LOG_E(context, "The shape of softmaxSum is invalid, got (%ld,%ld,%ld), should be (%u,%u,8)", dim0, dim1,
                  dim2, t1, n1);
        return ge::GRAPH_FAILED;
    }
    return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckAttentionInShape(gert::TilingContext *context)
{
    auto attentionInShape = context->GetOptionalInputShape(ATTENTION_IN);
    if (attentionInShape == nullptr) {
        return ge::GRAPH_SUCCESS;
    }
    auto queryShape = context->GetInputShape(QUERY);
    auto attentionInShapeDim = attentionInShape->GetStorageShape().GetDimNum();
    auto queryShapeDim = queryShape->GetStorageShape().GetDimNum();
    if (attentionInShapeDim != queryShapeDim) {
        OPS_LOG_E(context, "The dimnum of attentionIn %zu should be equal to query %zu", attentionInShapeDim,
                  queryShapeDim);
        return ge::GRAPH_FAILED;
    }
    for (size_t i = 0; i < queryShapeDim; i++) {
        if (attentionInShape->GetStorageShape().GetDim(i) != queryShape->GetStorageShape().GetDim(i)) {
            OPS_LOG_E(context, "The dim %zu of attentionIn shape is invalid, got %ld, should be %ld", i,
                      attentionInShape->GetStorageShape().GetDim(i), queryShape->GetStorageShape().GetDim(i));
            return ge::GRAPH_FAILED;
        }
    }
    return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckShapeValid(gert::TilingContext *context, uint32_t b, uint32_t n1, uint32_t s1, uint32_t d)
{
    auto isShapeInValid = (b == 0 || n1 == 0 || s1 == 0 || d == 0);
    if (isShapeInValid) {
        OPS_LOG_E(context, "input shape error, got 0 in bnsd(%u,%u,%u,%u)", b, n1, s1, d);
        return ge::GRAPH_FAILED;
    }

    auto ret = CheckSoftmaxMaxShape(context, b, n1, s1);
    if (ret != ge::GRAPH_SUCCESS) {
        return ret;
    }
    ret = CheckSoftmaxSumShape(context, b, n1, s1);
    if (ret != ge::GRAPH_SUCCESS) {
        return ret;
    }
    ret = CheckAttentionInShape(context);
    if (ret != ge::GRAPH_SUCCESS) {
        return ret;
    }

    return ge::GRAPH_SUCCESS;
}

ge::graphStatus CheckTndShapeValid(gert::TilingContext *context, uint32_t t1, uint32_t n1, uint32_t d)
{
    if (context == nullptr) {
        OPS_LOG_E("CheckTndShapeValid", "context is nullptr");
        return ge::GRAPH_FAILED;
    }

    auto isShapeInValid = (t1 == 0 || n1 == 0 || d == 0);
    if (isShapeInValid) {
        OPS_LOG_E(context, "input shape error, got 0 in tnd(%u,%u,%u)", t1, n1, d);
        return ge::GRAPH_FAILED;
    }

    auto ret = CheckTndSoftmaxMaxShape(context, t1, n1);
    if (ret != ge::GRAPH_SUCCESS) {
        return ret;
    }
    ret = CheckTndSoftmaxSumShape(context, t1, n1);
    if (ret != ge::GRAPH_SUCCESS) {
        return ret;
    }
    ret = CheckAttentionInShape(context);
    if (ret != ge::GRAPH_SUCCESS) {
        return ret;
    }

    return ge::GRAPH_SUCCESS;
}

} // namespace optiling
